
""" Siamese implementation using Tensorflow with MNIST example.
This siamese network embeds a 28x28 image (a point in 784D)
into a point in 2D.

By Youngwook Paul Kwon (young at berkeley.edu)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np
import data5.config as config
import data5.load_image as load_image
import data5.preprocessing as preprocessing
import cv2
import data5.inference as inference

# prepare data and tf.session
sess = tf.InteractiveSession()

# setup siamese network
siamese = inference.siamese()
train_step = tf.train.GradientDescentOptimizer(0.001).minimize(siamese.loss)
saver = tf.train.Saver()
tf.global_variables_initializer().run()


image_label_gen_batch9 = load_image.generator_next_batch(config.yzm_path, config.suffix, preprocessing.mapping_line_parser,
                                                 preprocessing.preprocessing, cut_image_func=preprocessing.cut_image, batch_size=9)

next_batch = load_image.generator_next_batch(config.yzm_path, config.suffix, preprocessing.mapping_line_parser,
                                                 preprocessing.preprocessing, cut_image_func=preprocessing.cut_image)



# start training
for step in range(200000):
    imgs, items, labels = next(next_batch)
    feed_dict = {
        siamese.x1: np.reshape(imgs, [50, config.char_image_height * config.char_image_width]),
        siamese.x2: np.reshape(items, [50, config.char_image_height * config.char_image_width]),
        siamese.y_: labels.astype('float')}
    _, loss_v = sess.run([train_step, siamese.loss], feed_dict)


    if np.isnan(loss_v):
        print('Model diverged with loss = NaN')
        quit()

    if step > 0 and step % 200 == 0:
        print ('step %d: loss %.3f' % (step, loss_v))
        # ===========================
        true_num = 0
        for j in range(10):

            imgs_9_1, imgs_9_2, labels_9 = next(image_label_gen_batch9)
            eucd = sess.run('eucd:0', feed_dict={
                siamese.x1: np.reshape(imgs_9_1, [9, 45 * 45]),
                siamese.x2: np.reshape(imgs_9_2, [9, 45 * 45]),
            })

            index = np.argmin(eucd)  # 预测的index
            index_ = np.where(labels_9 == True)[0][0]  # 正确的index
            if (index == index_):
                true_num += 1
        print(true_num / 10)

    # if step % 1000 == 0 and step > 0:
saver.save(sess, './ckpt/model')

